import asyncio
import gzip
import json
import uuid
import websockets
from settings import settings

PROTOCOL_VERSION = 0b0001
DEFAULT_HEADER_SIZE = 0b0001
FULL_CLIENT_REQUEST = 0b0001
AUDIO_ONLY_REQUEST = 0b0010
FULL_SERVER_RESPONSE = 0b1001
SERVER_ACK = 0b1011
SERVER_ERROR_RESPONSE = 0b1111
NO_SEQUENCE = 0b0000
POS_SEQUENCE = 0b0001
NEG_SEQUENCE = 0b0010
NEG_WITH_SEQUENCE = 0b0011
JSON_SERIAL = 0b0001
NO_SERIALIZATION = 0b0000
GZIP = 0b0001
NO_COMPRESSION = 0b0000


def generate_header(message_type=FULL_CLIENT_REQUEST, message_type_specific_flags=NO_SEQUENCE, serial_method=JSON_SERIAL, compression_type=GZIP, reserved_data=0x00):
    header = bytearray()
    header_size = 1
    header.append((PROTOCOL_VERSION << 4) | header_size)
    header.append((message_type << 4) | message_type_specific_flags)
    header.append((serial_method << 4) | compression_type)
    header.append(reserved_data)
    return header

def generate_before_payload(sequence: int):
    before_payload = bytearray()
    before_payload.extend(sequence.to_bytes(4, 'big', signed=True))
    return before_payload

def parse_response(res):
    protocol_version = res[0] >> 4
    header_size = res[0] & 0x0f
    message_type = res[1] >> 4
    message_type_specific_flags = res[1] & 0x0f
    serialization_method = res[2] >> 4
    message_compression = res[2] & 0x0f
    reserved = res[3]
    header_extensions = res[4:header_size * 4]
    payload = res[header_size * 4:]
    result = {'is_last_package': False}
    payload_msg = None
    payload_size = 0
    if message_type_specific_flags & 0x01:
        seq = int.from_bytes(payload[:4], "big", signed=True)
        result['payload_sequence'] = seq
        payload = payload[4:]
    if message_type_specific_flags & 0x02:
        result['is_last_package'] = True
    if message_type == FULL_SERVER_RESPONSE:
        payload_size = int.from_bytes(payload[:4], "big", signed=True)
        payload_msg = payload[4:]
    elif message_type == SERVER_ACK:
        seq = int.from_bytes(payload[:4], "big", signed=True)
        result['seq'] = seq
        if len(payload) >= 8:
            payload_size = int.from_bytes(payload[4:8], "big", signed=False)
            payload_msg = payload[8:]
    elif message_type == SERVER_ERROR_RESPONSE:
        code = int.from_bytes(payload[:4], "big", signed=False)
        result['code'] = code
        payload_size = int.from_bytes(payload[4:8], "big", signed=False)
        payload_msg = payload[8:]
    if payload_msg is None:
        return result
    if message_compression == GZIP:
        payload_msg = gzip.decompress(payload_msg)
    if serialization_method == JSON_SERIAL:
        payload_msg = json.loads(str(payload_msg, "utf-8"))
    elif serialization_method != NO_SERIALIZATION:
        payload_msg = str(payload_msg, "utf-8")
    result['payload_msg'] = payload_msg
    result['payload_size'] = payload_size
    return result

class DoubaoASRClient:
    def __init__(self, on_result_callback=None):
        self.ws_url = settings.DOUBAO_WS_URL
        self.app_key = settings.DOUBAO_APP_KEY
        self.access_key = settings.DOUBAO_ACCESS_KEY
        self.resource_id = settings.DOUBAO_RESOURCE_ID
        self.on_result_callback = on_result_callback
        self.ws = None
        self.seq = 1
        self.reqid = str(uuid.uuid4())
        self.connect_id = str(uuid.uuid4())  # 用于追踪连接的ID
        self.connected = False
        self._send_lock = asyncio.Lock()

    async def connect(self):
        header = {
            "X-Api-Resource-Id": self.resource_id,
            "X-Api-Access-Key": self.access_key,
            "X-Api-App-Key": self.app_key,
            "X-Api-Request-Id": self.reqid,
            "X-Api-Connect-Id": self.connect_id  # 必须的追踪ID
        }
        
        print(f"[豆包ASR] 连接参数:")
        print(f"  URL: {self.ws_url}")
        print(f"  AppKey: {self.app_key[:8]}...")
        print(f"  AccessKey: {self.access_key[:8]}...")
        print(f"  ResourceId: {self.resource_id}")
        print(f"  RequestId: {self.reqid}")
        print(f"  ConnectId: {self.connect_id}")
        
        try:
            print(f"[豆包ASR] 正在连接豆包API...")
            self.ws = await websockets.connect(self.ws_url, additional_headers=header, max_size=1000000000)
            print(f"[豆包ASR] WebSocket连接已建立")
            
            # 发送Full client request
            request_params = {
                "user": {"uid": self.reqid},
                "audio": {
                    'format': "pcm",  # 明确指定为PCM格式
                    "sample_rate": 16000,
                    "bits": 16,
                    "channel": 1,
                    "codec": "raw",
                },
                "request": {
                    "model_name": "bigmodel",
                    "enable_punc": True,
                    "result_type": "single", #增加返回识别结果
                }
            }
            
            print(f"[豆包ASR] 发送请求参数: {json.dumps(request_params, indent=2)}")
            
            payload_bytes = str.encode(json.dumps(request_params))
            payload_bytes = gzip.compress(payload_bytes)
            full_client_request = bytearray(generate_header(message_type_specific_flags=POS_SEQUENCE))
            full_client_request.extend(generate_before_payload(sequence=self.seq))
            full_client_request.extend((len(payload_bytes)).to_bytes(4, 'big'))
            full_client_request.extend(payload_bytes)
            
            print(f"[豆包ASR] 发送Full client request，长度: {len(full_client_request)} 字节")
            await self.ws.send(full_client_request)
            
            print(f"[豆包ASR] 等待服务器响应...")
            res = await self.ws.recv()
            print(f"[豆包ASR] 收到响应，长度: {len(res)} 字节")
            
            result = parse_response(res)
            print(f"[豆包ASR] 解析响应: {json.dumps(result, indent=2, default=str)}")
            
            # 检查连接是否成功
            if 'payload_msg' in result and isinstance(result['payload_msg'], dict):
                payload_msg = result['payload_msg']
                print(f"[豆包ASR] 响应payload: {json.dumps(payload_msg, indent=2, default=str)}")
                
                # 检查是否有错误码
                if 'code' in payload_msg:
                    if payload_msg.get('code') != 1000:  # 豆包API成功码是1000
                        error_msg = payload_msg.get('message', '未知错误')
                        error_code = payload_msg.get('code', '未知错误码')
                        print(f"[豆包ASR] 豆包API返回错误: code={error_code}, message={error_msg}")
                        raise Exception(f"豆包API连接失败: code={error_code}, message={error_msg}")
                    else:
                        print(f"[豆包ASR] 豆包API连接成功")
                else:
                    # 没有错误码，检查是否有错误信息
                    if 'error' in payload_msg or 'message' in payload_msg:
                        error_msg = payload_msg.get('error', payload_msg.get('message', '未知错误'))
                        print(f"[豆包ASR] 豆包API返回错误: {error_msg}")
                        raise Exception(f"豆包API连接失败: {error_msg}")
                    else:
                        # 没有错误码也没有错误信息，认为是成功响应
                        print(f"[豆包ASR] 豆包API连接成功（初始化响应）")
            else:
                print(f"[豆包ASR] 响应格式异常，没有payload_msg")
                raise Exception(f"豆包API响应格式异常: {result}")
            
            self.seq += 1
            self.connected = True
            print(f"[豆包ASR] 连接成功，Connect-Id: {self.connect_id}")
            # 启动后台接收任务
            asyncio.create_task(self._recv_loop())
        except Exception as e:
            print(f"[豆包ASR] 连接失败: {e}")
            if hasattr(e, '__traceback__'):
                import traceback
                print(f"[豆包ASR] 详细错误信息:")
                traceback.print_exc()
            raise

    async def send_audio(self, chunk: bytes, is_last=False):
        async with self._send_lock:
            seq = self.seq
            if is_last:
                seq = -seq
            payload_bytes = gzip.compress(chunk)
            if is_last:
                audio_only_request = bytearray(generate_header(message_type=AUDIO_ONLY_REQUEST, message_type_specific_flags=NEG_WITH_SEQUENCE))
            else:
                audio_only_request = bytearray(generate_header(message_type=AUDIO_ONLY_REQUEST, message_type_specific_flags=POS_SEQUENCE))
            audio_only_request.extend(generate_before_payload(sequence=seq))
            audio_only_request.extend((len(payload_bytes)).to_bytes(4, 'big'))
            audio_only_request.extend(payload_bytes)
            await self.ws.send(audio_only_request)
            self.seq += 1

    async def _recv_loop(self):
        try:
            while True:
                res = await self.ws.recv()
                result = parse_response(res)
                
                # 处理错误响应
                if result.get('code') and result['code'] != 1000:
                    print(f"[豆包ASR] 错误码: {result['code']}")
                    break
                
                if 'payload_msg' in result and self.on_result_callback:
                    text = self._extract_text(result['payload_msg'])
                    if text:
                        print(f"[豆包ASR] 收到识别结果: '{text}'")
                        self.on_result_callback(text)
                if result.get('is_last_package'):
                    break
        except Exception as e:
            print(f"[豆包ASR] 接收异常: {e}")

    def _extract_text(self, payload_msg):
        # 解析豆包API返回的识别文本
        try:
            if isinstance(payload_msg, dict):
                if 'result' in payload_msg and 'text' in payload_msg['result']:
                    return payload_msg['result']['text']
                elif 'result' in payload_msg and isinstance(payload_msg['result'], list):
                    # 兼容多段
                    return ''.join([seg.get('text', '') for seg in payload_msg['result']])
        except Exception as e:
            print(f"[豆包ASR] 文本提取异常: {e}")
        return None

    async def close(self):
        if self.ws:
            await self.ws.close()
        self.connected = False 